# coding: utf-8
#
import argparse
import re
import os
from re import escape
import time
import hashlib
import glob
import json
import pprint
import shutil
import socket
import subprocess
import sys
import tempfile
import multiprocessing
import uiautomator2 as u2
import requests
import tornado.web
import settings
import adbutils
import apkutils2 as apkutils

from concurrent.futures import ThreadPoolExecutor
from uiautomator2 import Device
from functools import partial
from logzero import logger

from tornado import gen, websocket
from tornado.concurrent import run_on_executor
from tornado.ioloop import IOLoop
from tornado.web import RequestHandler
from tornado.websocket import WebSocketHandler, websocket_connect

from device import STATUS_FAIL, STATUS_INIT, STATUS_OKAY, AndroidDevice
from asyncadb import adb
from adbutils import adb as adbclient
from core.utils import current_ip, fix_url, id_generator, update_recursive
from core import fetching
from heartbeat import heartbeat_connect
import datetime
from remote.wechat.operators import App
from remote.jobs import jobManager, Job
import remote.settings as rs
from remote.settings import server_setting

__curdir__ = os.path.dirname(os.path.abspath(__file__))

hbconn = None
udid2device = {}
secret = id_generator(10)
PROVIDER_TASK_TYPE = ''

class CorsMixin(object):
    CORS_ORIGIN = '*'
    CORS_METHODS = 'GET,POST,OPTIONS'
    CORS_CREDENTIALS = True
    CORS_HEADERS = "x-requested-with,authorization"

    def set_default_headers(self):
        self.set_header("Access-Control-Allow-Origin", self.CORS_ORIGIN)
        self.set_header("Access-Control-Allow-Headers", self.CORS_HEADERS)
        self.set_header('Access-Control-Allow-Methods', self.CORS_METHODS)
        self.set_header('Access-Control-Allow-Headers', 'Content-Type')
    def options(self):
        # no body
        self.set_status(204)
        self.finish()

class AppHandler(CorsMixin, tornado.web.RequestHandler):
    _install_executor = ThreadPoolExecutor(4)
    _download_executor = ThreadPoolExecutor(1)

    def cache_filepath(self, text: str) -> str:
        m = hashlib.md5()
        m.update(text.encode('utf-8'))
        return "cache-" + m.hexdigest()

    @run_on_executor(executor="_download_executor")
    def cache_download(self, url: str) -> str:
        """ download with local cache """
        target_path = self.cache_filepath(url)
        logger.debug("Download %s to %s", url, target_path)

        if os.path.exists(target_path):
            logger.debug("Cache hited")
            return target_path

        # TODO: remove last
        for fname in glob.glob("cache-*"):
            logger.debug("Remove old cache: %s", fname)
            os.unlink(fname)

        tmp_path = target_path + ".tmp"
        r = requests.get(url, stream=True)
        r.raise_for_status()

        with open(tmp_path, "wb") as tfile:
            content_length = int(r.headers.get("content-length", 0))
            if content_length:
                for chunk in r.iter_content(chunk_size=40960):
                    tfile.write(chunk)
            else:
                shutil.copyfileobj(r.raw, tfile)

        os.rename(tmp_path, target_path)
        return target_path

    @run_on_executor(executor='_install_executor')
    def app_install_url(self, serial: str, apk_path: str, **kwargs):
        pkg_name = app_install_local(serial, apk_path, **kwargs)
        return {
            "success": True,
            "description": "Success",
            "packageName": pkg_name,
        }

    async def post(self, udid=None):
        udid = udid or self.get_argument("udid")
        device = udid2device[udid]
        url = self.get_argument("url")
        launch = self.get_argument("launch",
                                   "false") in ['true', 'True', 'TRUE', '1']

        try:
            apk_path = await self.cache_download(url)
            ret = await self.app_install_url(device.serial,
                                             apk_path,
                                             launch=launch)
            self.write(ret)
        except InstallError as e:
            self.set_status(400)
            self.write({
                "success": False,
                "description": "{}: {}".format(e.stage, e.reason)
            })
        except Exception as e:
            self.set_status(500)
            self.write(str(e))

class ColdingHandler(tornado.web.RequestHandler):
    async def post(self, udid=None):
        """ 设备清理 """
        udid = udid or self.get_argument("udid")
        logger.info("Receive colding request for %s", udid)
        request_secret = self.get_argument("secret")
        if secret != request_secret:
            logger.warning("secret not match, expect %s, got %s", secret,
                           request_secret)
            return

        if udid not in udid2device:
            return

        device = udid2device[udid]
        await device.reset()
        await hbconn.device_update({
            "udid": udid,
            "colding": False,
            "provider": device.addrs(),
        })
        self.write({"success": True, "description": "Device colded"})

class InstallError(Exception):
    def __init__(self, stage: str, reason):
        self.stage = stage
        self.reason = reason

class MobileJobHandler(CorsMixin, tornado.web.RequestHandler):
    '''处理服务端下发的任务'''
    async def post(self, *args, **kwargs):
        """ 分配wechat工作任务 """
        data = json.loads(self.request.body)
        udid = data["udid"] 

        logger.info("Receive task request for %s", udid)
        # validate secret
        request_secret = self.get_argument("secret")
        if secret != request_secret:
            logger.warning("secret not match, expect %s, got %s", secret,
                           request_secret)
            return

        # check if the device is online.
        if udid not in udid2device:
            return
            
        device = udid2device[udid]
        job_id = data["job_id"]
        app_name = data["app_name"]
        job_type = data["job_type"]
        task_load = int(data["task_load"])
        stri = data["plan_datetime"]
        plan_datetime = datetime.datetime.strptime(stri, "%Y-%m-%d %H:%M:%S")
        operator_data = data["operator_data"]
        task_generating_pattern = data["task_generating_pattern"]
        
        job = Job(job_id, device.serial, app_name, job_type, task_load, plan_datetime, operator_data, task_generating_pattern)
        
        # 如果前序任务没有完全执行完毕，接到新的任务，移除前序任务
        await jobManager.clear_jobs(job_id)
        
        jobManager.try_accept(job)
        logger.info("任务 %s 的任务已经被接受。", job_id)
        # await hbconn.job_update({
        #     "udid": udid,
        #     "completedtaskLoad": False,
        #     "provider": device.addrs(),
        # })
        success = True
        msg = "任务 " + job_id + " 的任务已经被接受。"
        self.write({"success": success, "description": msg})

def make_app():
    """
    配置 REST 服务路由。
    """
    app = tornado.web.Application([
        (r"/app/install", AppHandler),
        (r"/app/mobile/job", MobileJobHandler),
        (r"/cold", ColdingHandler)
    ])
    return app

async def device_watch(allow_remote: bool = False, task_type: str = None):
    serial2udid = {}
    udid2serial = {}

    def callback(udid: str, status: str):
        if status == STATUS_OKAY:
            print("Good")

    async for event in adb.track_devices():
        logger.debug("%s", event)
        # udid = event.serial  
        # FIXME (ssx): fix later

        if not allow_remote:
            if re.match(r"(\d+)\.(\d+)\.(\d+)\.(\d+):(\d+)", event.serial):
                logger.debug("Skip remote device: %s", event)
                continue
        if event.present:
            try:
                udid = serial2udid[event.serial] = event.serial
                udid2serial[udid] = event.serial

                device = AndroidDevice(event.serial, partial(callback, udid))

                await device.init()
                await device.open_identify()

                udid2device[udid] = device
                
                await hbconn.device_update({
                    # "private": False, # TODO
                    "udid": udid,
                    "platform": "android",
                    "colding": False,
                    "provider": device.addrs(),
                    "properties": await device.properties(),
                    "providerTaskType": task_type,
                    "appName": "wechat" #目前主要对象是wechat
                }) # yapf: disable
                
                logger.info("Device:%s is ready", event.serial)
            except RuntimeError:
                logger.warning("Device:%s initialize failed", event.serial)
            except Exception as e:
                logger.error("Unknown error: %s", e)
                import traceback
                traceback.print_exc()
        else:
            udid = serial2udid[event.serial]
            #添加清理wechat task的逻辑
            if udid in udid2device:
                udid2device[udid].close()
                udid2device.pop(udid, None)

            await hbconn.device_update({
                "udid": udid,
                "provider": None,  # not present
            })

def app_install_local(serial: str, apk_path: str, launch: bool = False) -> str:
    """
    install apk to device

    Returns:
        package name

    Raises:
        AdbInstallError, FileNotFoundError
    """
    # 解析apk文件
    device = adbclient.device(serial)
    try:
        apk = apkutils.APK(apk_path)
    except apkutils.apkfile.BadZipFile:
        raise InstallError("ApkParse", "Bad zip file")

    # 提前将重名包卸载
    package_name = apk.manifest.package_name
    pkginfo = device.package_info(package_name)
    if pkginfo:
        logger.debug("uninstall: %s", package_name)
        device.uninstall(package_name)

    # 解锁手机，防止锁屏
    # ud = u2.connect_usb(serial)
    # ud.open_identify()
    try:
        # 推送到手机
        dst = "/data/local/tmp/tmp-%d.apk" % int(time.time() * 1000)
        logger.debug("push %s %s", apk_path, dst)
        device.sync.push(apk_path, dst)
        logger.debug("install-remote %s", dst)
        # 调用pm install安装
        device.install_remote(dst)
    except adbutils.errors.AdbInstallError as e:
        raise InstallError("install", e.output)
    # finally:
    # 停止uiautomator2服务
    # logger.debug("uiautomator2 stop")
    # ud.session().press("home")
    # ud.service("uiautomator").stop()

    # 启动应用
    if launch:
        logger.debug("launch %s", package_name)
        device.app_start(package_name)
    return package_name

async def async_main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    # yapf: disable
    parser.add_argument('-s', '--server', default='localhost:4000', help='server address')
    parser.add_argument("--allow-remote", action="store_true", help="allow remote connect device")
    parser.add_argument('-t', '--test', default=False, type=bool, help="run test code")
    parser.add_argument('--provider-url', type=str, default='http://localhost', help='provider address')
    parser.add_argument('-p', '--port', type=int, default=3500, help='listen port')
    parser.add_argument("--atx-agent-version", default=u2.version.__atx_agent_version__, help="set atx-agent version")
    parser.add_argument("--owner", type=str, help="provider owner email")
    parser.add_argument("--owner-file", type=argparse.FileType("r"), help="provider owner email from file")
    parser.add_argument("--task-type", default='adfe', type=str, help="provider the type of task: adfg, adfe")
    args = parser.parse_args()
    settings.ATX_AGENT_VERSION = args.atx_agent_version

    owner_email = args.owner
    if args.owner_file:
        with args.owner_file as file:
            owner_email = file.read().strip()
    logger.info("Owner: %s", owner_email)
    
    if args.task_type != "adfg" and args.task_type != "adfe":
        raise RuntimeError("--task-type just can accept neither 'adfg' or 'adfe'.")
    
    task_type = args.task_type
    logger.info("task-type: %s", args.task_type)

    # if args.test:
    #     for apk_name in ("cloudmusic.apk", ):  # , "apkinfo.exe"):
    #         apk_path = "testdata/" + apk_name
    #         logger.info("Install %s", apk_path)
    #         # apk_path = r"testdata/cloudmusic.apk"
    #         ret = app_install_local("6EB0217704000486", apk_path, launch=True)
    #         logger.info("Ret: %s", ret)
    #     return

    # start local server
    # provider_url = "http://" + current_ip() + ":" + str(args.port)
    
    provider_url = args.provider_url + ":" + str(args.port)
    app = make_app()
    app.listen(args.port)
    logger.info("ProviderURL: %s", provider_url)

    fetching.get_all()

    
    server_setting.server_url = "http://" + args.server
    server_setting.test = args.test
    # connect to atxserver2
    global hbconn
    hbconn = await heartbeat_connect(args.server,
                                     secret=secret,
                                     self_url=provider_url,
                                     owner=owner_email)

    await device_watch(args.allow_remote, task_type)

if __name__ == "__main__":
    try:
        IOLoop.current().run_sync(async_main)
    except KeyboardInterrupt:
        logger.info("Interrupt catched")